/**
 * @file pcap_monitor.cpp
 * @author your name (you@domain.com)
 * @brief 监听指定网络访问
 * @version 0.1
 * @date 2023-05-19
 *
 * @note
 *     1、简单去重，同一个地址，1S内只允许有一次记录
 *
 * @copyright Copyright (c) 2023
 *
 */

#include "pcap_monitor.h"

#include <cstring>
#include <iostream>
#include <map>
#include <string>
#include <thread>
#include <chrono>

#include <stdio.h>
#include <unistd.h>
#include <netdb.h>
#include <pcap.h>
#include <fcntl.h>
#include <arpa/inet.h>
#include <net/ethernet.h>
#include <netinet/ip.h>
#include <netinet/tcp.h>
#include <netinet/in.h>

#define IPV4 4
int g_timeSpaces = 1; // 去重间隔时间S

void run(struct pcap *, void *);
void packet_handler(u_char *, const struct pcap_pkthdr *, const u_char *);

PcapMonitor::PcapMonitor(QObject *parent)
    : QObject(parent)
    , m_pcapHandle(nullptr)
    , m_thread(nullptr)
    , m_isRunning(false)
{
}

PcapMonitor::~PcapMonitor()
{
    if (m_pcapHandle) {
        stop();
    }
}

std::vector<std::string> PcapMonitor::getAllNetInterfaceNames()
{
    std::vector<std::string> vInterfaces;
    char errbuf[PCAP_ERRBUF_SIZE];
    pcap_if_t *interfaces;

    if (pcap_findalldevs(&interfaces, errbuf) == -1) {
        fprintf(stderr, "Error finding devices: %s\n", errbuf);
        return vInterfaces;
    }

    for (pcap_if_t *interface = interfaces; interface != NULL; interface = interface->next) {
        // Warning:这里去掉部分网卡涉笔，否则后续抓包会崩溃
        if ("nflog" == std::string(interface->name)) {
            continue;
        }
        vInterfaces.push_back(interface->name);
    }

    pcap_freealldevs(interfaces);
    return vInterfaces;
}

/**
 * @brief 生成抓包过滤规则，目前只抓包tcp
 * @param ips
 * @param ports
 * @return std::string
 */
std::string PcapMonitor::makeMonitorRule(std::list<std::string> ips, std::list<std::string> ports)
{
    std::string rule = "tcp";
    if (ips.size() > 0) {
        rule += " and dst host (";
        for (auto ite = ips.begin(); ite != ips.end();) {
            rule += *ite;
            ite++;
            if (ite != ips.end()) {
                rule += " or ";
            }
        }
        rule += ")";
    }

    if (ports.size() > 0) {
        rule += " and dst port (";
        for (auto ite = ports.begin(); ite != ports.end();) {
            rule += *ite;
            ite++;
            if (ite != ports.end()) {
                rule += " or ";
            }
        }
        rule += ")";
    }

    return rule;
}

/**
 * @brief 设置抓包过滤规则
 * @param rule
 */
void PcapMonitor::setMonitorRule(const std::string &rule)
{
    m_rule = rule;
}

bool PcapMonitor::start(const std::string &interface)
{
    // 已经启动
    if (m_pcapHandle) {
        return false;
    }

    // 数据异常
    if (!m_rule.length() || !interface.length()) {
        return false;
    }

    char errbuf[PCAP_ERRBUF_SIZE];
    // 打开网卡，设置混杂模式和过滤规则（捕获TCP端口为80或443的流量）
    m_pcapHandle = pcap_open_live(interface.c_str(), BUFSIZ, 0, 1000, errbuf);
    if (!m_pcapHandle) {
        std::cout << "pcap_open_live open [" << interface << "] failed!" << std::endl;
        return false;
    }
    // 设置非阻塞调用
    pcap_setnonblock(m_pcapHandle, 1, errbuf);

    struct bpf_program fp;
    bpf_u_int32 net;
    pcap_compile(m_pcapHandle, &fp, m_rule.c_str(), 1, net);
    pcap_setfilter(m_pcapHandle, &fp);

    m_thread = new std::thread(run, m_pcapHandle, (void *)this);
    return true;
}

void PcapMonitor::stop()
{
    if (m_pcapHandle) {
        m_isRunning.store(false);
        m_thread->join();

        m_pcapHandle = nullptr;
        delete m_thread;
        m_thread = nullptr;
    }
}

void PcapMonitor::notify(const std::string &ip, const uint16_t port)
{
    // Warning:这里回调通知跟类对象不在一个线程，注意数据竞态问题{m_records}
    auto current_time = std::chrono::high_resolution_clock::now();
    auto duration_in_ms = std::chrono::duration_cast<std::chrono::milliseconds>(current_time.time_since_epoch());
    long ms = duration_in_ms.count();
    RecordKey key {ip, port};

    // 没有记录
    if (m_records.find(key) == m_records.end()) {
        m_records[key] = ms;
        std::cout << "IP:" << ip << " Port:" << port << std::endl;
        accessNotify(ip.c_str(), port);
        return;
    }

    auto lastTimeMs = m_records[key];
    m_records[key] = ms;
    // 距离上次记录超过{g_timeSpaces}S
    if ((ms - lastTimeMs) > g_timeSpaces * 1000) {
        std::cout << "IP:" << ip << " Port:" << port << std::endl;
        accessNotify(ip.c_str(), port);
    }
}

/**
 * @brief 线程主体函数
 * @param pcapHandle
 * @param userData PcapMonitor对象
 */
void run(struct pcap *pcapHandle, void *userData)
{
    auto pcapMonitor = (PcapMonitor *)userData;
    pcapMonitor->m_isRunning.store(true);
    while (pcapMonitor->m_isRunning.load()) {
        struct pcap_pkthdr *header;
        const u_char *packet_data;

        int res = pcap_next_ex(pcapHandle, &header, &packet_data);
        if (res == 1) {
            // 成功获取到一个数据包
            // 处理抓取到的数据包内容
            packet_handler((u_char *)userData, header, packet_data);
            continue;
        } else if (res == -1) {
            // 发生错误
            break;
        }
        usleep(10 * 1000);
    }
    pcap_close(pcapHandle);
    pcapMonitor->m_isRunning.store(false);
}

/**
 * @brief 回调函数，处理捕获到的数据包
 * @note 目前只处理了IPV4 TCP数据
 * @param args PcapMonitor对象
 * @param header 包信息
 * @param packet 数据包
 */
void packet_handler(u_char *args, const struct pcap_pkthdr *header, const u_char *packet)
{
    struct iphdr *ip_hdr = (struct iphdr *)(packet + sizeof(struct ether_header));
    char src_ip[INET_ADDRSTRLEN];
    char dst_ip[INET_ADDRSTRLEN];
    inet_ntop(AF_INET, &ip_hdr->saddr, src_ip, INET_ADDRSTRLEN);
    inet_ntop(AF_INET, &ip_hdr->daddr, dst_ip, INET_ADDRSTRLEN);

    // 判断是否是IPV4协议
    if (ip_hdr->version != IPV4) {
        return;
    }
    // 判断是否是TCP协议
    if (ip_hdr->protocol != IPPROTO_TCP) {
        return;
    }

    struct tcphdr *tcp_hdr = (struct tcphdr *)(packet + sizeof(struct ether_header) + sizeof(struct iphdr));
    if (args) {
        auto pcapMonitor = (PcapMonitor *)args;
        pcapMonitor->notify(dst_ip, ntohs(tcp_hdr->th_dport));
    }
}
